package org.zjx.server;

import com.alibaba.fastjson2.JSON;
import lombok.Getter;
import lombok.SneakyThrows;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.zjx.client.SocketClient;
import org.zjx.core.SocketConfig;
import org.zjx.exception.RpcTimeoutException;
import org.zjx.exception.ServiceUnavailableException;
import org.zjx.handler.CommandHandler;
import org.zjx.handler.CommandHandlerFactory;
import org.zjx.message.IMessage;
import org.zjx.message.messages.*;
import org.zjx.server.session.ClientSession;
import org.zjx.server.session.SessionManager;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.lang.reflect.Method;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.net.Socket;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.stream.Collectors;

@Getter
public class ClusterHandler extends ClientHandler {
    private static final Logger logger = LoggerFactory.getLogger(ClusterHandler.class);
    private final ClusterServer server;
    private final ClientSession session;

    public ClusterHandler(ClusterServer server, Socket socket, DataInputStream dataIn, DataOutputStream dataOut) {
        super(server, socket, dataIn, dataOut);
        this.server = server;
        this.session = SessionManager.getInstance().createSession(socket, dataOut);
    }

    // 添加获取客户端地址的方法
    private String getClientAddress() {
        if (socket != null && socket.isConnected()) {
            return socket.getInetAddress().getHostAddress() + ":" + socket.getPort();
        }
        return "unknown";
    }

    @Override
    protected void cleanupConnection() {
        // 1. 标记连接不活跃
        connectionActive = false;

        // 2. 清理 SessionManager 中的会话
        if (session != null) {
            SessionManager.getInstance().removeSession(session.getSessionId());
            logger.info("Session removed: {}", session.getSessionId());
        }

        // 3. 关闭流和 Socket（调用父类方法）
        super.cleanupConnection();
    }


    @Override
    public void processMessage(IMessage message) {
        try {

            if (message instanceof RpcRequest) {
                handleRpcRequest((RpcRequest) message);
            } else if (message instanceof CommandMessage) {
                new Thread(() -> {
                    handleCommand((CommandMessage) message);
                }).start();
            } else if (message instanceof HeartbeatMessage) {
                handleHeartbeat((HeartbeatMessage) message);
            } else if (message instanceof RegisterWorkerMessage) {
                handleWorkerRegister((RegisterWorkerMessage) message);
            } else if (message instanceof RouteQueryMessage) {
                handleRouteQuery((RouteQueryMessage) message);
            } else {
                handleOtherMessage(message);
            }
        } catch (Exception e) {
            logger.error("Message processing error: " + message, e);
            sendErrorResponse("Message processing error");
        }
    }

    private void handleCommand(CommandMessage command) {
        // 特殊处理注册命令
        if ("REGISTER".equals(command.getCommand())) {
            handleRegistration(command);
            return;
        }

        ResponseMessage responseMessage = CommandHandlerFactory.getHandler(command.getCommand())
                .map(handler -> handler.handle(command))
                .orElse(ResponseMessage.error(command, "Unsupported command"));
        sendMessage(responseMessage);
    }

    private void handleRegistration(CommandMessage cmd) {
        String clientId = (String) cmd.getParams().get("clientId");
        if (clientId != null && !clientId.isEmpty()) {
            session.setClientId(clientId);
            SessionManager.getInstance().bindClientToSession(clientId, session.getSessionId());

            Optional<CommandHandler> commandHandler = CommandHandlerFactory.getHandler("REGISTER");
            if (commandHandler.isPresent()) {
                sendMessage(commandHandler.get().handle(cmd));
            } else {
                sendMessage(ResponseMessage.success(cmd, "Registered successfully"));
            }
        } else {
            sendMessage(ResponseMessage.error(cmd, "Invalid client ID"));
        }
    }

    private void handleOtherMessage(IMessage message) {
        sendErrorResponse("Unsupported message type: " + message.getClass().getSimpleName());
    }

    private void handleRpcRequest(RpcRequest request) {
        String serviceName = request.getClassName();
        logger.debug("Handling RPC request for service: {}", serviceName);

        if (server.getServiceRegistry().hasService(serviceName)) {
            processLocally(request);
            return;
        }

        RouteResponse route = getRouteForService(serviceName);

        if (route == null) {
            sendErrorResponse("Route service unavailable");
        } else if (route.isDirect()) {
            processLocally(request);
        } else if (route.isRedirect()) {
            proxyToWorker(request, route);
        } else if (route.isUnavailable()) {
            sendErrorResponse(route.getReason());
        }
    }

    private RouteResponse getRouteForService(String serviceName) {
        if (server.getServerType() == ClusterServer.ServerType.MASTER &&
                server.getRouteManager() != null) {

            return server.getRouteManager().queryRoute(serviceName);
        } else if (server.getServerType() == ClusterServer.ServerType.WORKER &&
                server.getMasterConnector() != null) {

            return server.getMasterConnector().queryRoute(serviceName);
        }
        return null;
    }

    @SneakyThrows
    private void processLocally(RpcRequest request) {
        Object service = server.getServiceRegistry().getService(request.getClassName());
        Class<?>[] paramTypes = request.getParameterTypes();
        Method method = findBestMatchingMethod(service.getClass(), request.getMethodName(), paramTypes);

        Object[] params = request.getParameters();
        logInvocationDetails(method, params);

        for (int i = 0; i < params.length; i++) {
            if (params[i] instanceof List) {
                Type paramType = method.getGenericParameterTypes()[i];
                if (paramType instanceof ParameterizedType) {
                    Class<?> elementType = (Class<?>) ((ParameterizedType) paramType)
                            .getActualTypeArguments()[0];
                    params[i] = convertList((List<?>) params[i], elementType);
                }
            }
        }

        Object result = method.invoke(service, params);

        RpcResponse response = new RpcResponse();
        response.setMessageId(request.getMessageId());
        response.setResult(result);
        sendMessage(response);
    }

    private <T> List<T> convertList(List<?> list, Class<T> targetType) {
        return list.stream()
                .map(item -> {
                    if (targetType.isInstance(item)) {
                        return targetType.cast(item);
                    } else {
                        return JSON.parseObject(JSON.toJSONString(item), targetType);
                    }
                })
                .collect(Collectors.toList());
    }

    /**
     * 查找最匹配的方法（处理重载情况）
     */
    private Method findBestMatchingMethod(Class<?> clazz, String methodName, Class<?>[] paramTypes)
            throws NoSuchMethodException {

        Method[] candidates = Arrays.stream(clazz.getMethods())
                .filter(m -> m.getName().equals(methodName))
                .filter(m -> m.getParameterCount() == paramTypes.length)
                .toArray(Method[]::new);

        if (candidates.length == 1) {
            return candidates[0];
        }

        // 寻找参数类型最匹配的方法
        for (Method method : candidates) {
            Class<?>[] methodParamTypes = method.getParameterTypes();
            boolean match = true;
            for (int i = 0; i < paramTypes.length; i++) {
                if (!methodParamTypes[i].isAssignableFrom(paramTypes[i])) {
                    match = false;
                    break;
                }
            }
            if (match) {
                return method;
            }
        }

        throw new NoSuchMethodException("No matching method found for " + methodName +
                " with parameters " + Arrays.toString(paramTypes));
    }

    private void logInvocationDetails(Method method, Object[] params) {
        logger.debug("Invoking method: {}.{}",
                method.getDeclaringClass().getSimpleName(),
                method.getName());
        logger.debug("Method parameter types: {}", Arrays.toString(method.getParameterTypes()));
        logger.debug("Actual parameter types: {}",
                Arrays.stream(params)
                        .map(p -> p != null ? p.getClass().getName() : "null")
                        .toArray());

        if (logger.isTraceEnabled()) {
            for (int i = 0; i < params.length; i++) {
                logger.trace("Parameter {} [{}]: {}",
                        i,
                        params[i] != null ? params[i].getClass().getName() : "null",
                        params[i]);
            }
        }
    }

    private void proxyToWorker(RpcRequest request, RouteResponse route) {
        logger.info("Proxying {} to worker {} at {}:{}",
                request.getClassName(),
                route.getWorkerId(),
                route.getTargetHost(),
                route.getTargetPort()
        );
        String workerId = route.getWorkerId();
        SocketClient workerClient = getWorkerConnection(route);

        try {
            IMessage response = workerClient.sendAndWait(request, server.getConfig().getRequestTimeout());
            if (response instanceof RpcResponse) {
                sendMessage(response);
            } else {
                sendErrorResponse("Invalid response from worker");
            }
        } catch (RpcTimeoutException e) {
            logger.error("Worker proxy timeout: {}", workerId, e);
            server.getWorkerConnections().remove(workerId);
            sendErrorResponse("Worker timeout");
        } catch (Exception e) {
            logger.error("Worker proxy failed: {}", workerId, e);
            server.getWorkerConnections().remove(workerId);
            sendErrorResponse("Worker connection error");
        }
    }

    private SocketClient getWorkerConnection(RouteResponse route) {
        return server.getWorkerConnections().computeIfAbsent(route.getWorkerId(), k -> {
            // 创建专用配置
            SocketConfig workerConfig = new SocketConfig();

            // 明确设置目标地址和端口
            workerConfig.setServerHost(route.getTargetHost());
            workerConfig.setPort(route.getTargetPort());

            // 禁用重连等高级功能（转发需要即时失败）
            workerConfig.setReconnectInterval(0);
            workerConfig.setMaxReconnectAttempts(0);

            // 设置合理的超时
            workerConfig.setRequestTimeout(server.getConfig().getRequestTimeout());

            // 禁用不必要的高级功能
            workerConfig.setReconnectInterval(0);
            workerConfig.setHeartbeatInterval(0);

            // 添加调试日志
            logger.info("Creating worker connection to {}:{}", workerConfig.getServerHost(), workerConfig.getPort());

            SocketClient client = new SocketClient(workerConfig);
            client.connect();
            return client;
        });
    }

    private void sendErrorResponse(String message) {
        RpcResponse response = new RpcResponse();
        response.setError(new ServiceUnavailableException(message));
        sendMessage(response);
    }

    private void handleWorkerRegister(RegisterWorkerMessage message) {
        if (server.getServerType() != ClusterServer.ServerType.MASTER) {
            return;
        }

        if (message.getWorkerId() == null) {
            logger.error("Received worker registration without workerId from {}", getClientAddress());
            sendErrorResponse("Missing workerId");
            return;
        }

        RouteManager.WorkerInfo worker = new RouteManager.WorkerInfo(
                message.getWorkerId(),
                message.getHost(),
                message.getPort(),
                System.currentTimeMillis(),
                message.getServices()
        );

        server.getRouteManager().registerWorker(worker);
        logger.info("Worker registered: {}", message.getWorkerId());

        AckMessage ackMessage = new AckMessage();
        ackMessage.setMessageId(message.getMessageId());

        sendMessage(ackMessage);
    }

    private void handleRouteQuery(RouteQueryMessage message) {
        if (server.getServerType() != ClusterServer.ServerType.MASTER ||
                server.getRouteManager() == null) {

            sendMessage(RouteResponse.unavailable(
                    message.getServiceName(), "Not a master server"));
            return;
        }

        RouteResponse response = server.getRouteManager().queryRoute(message.getServiceName());
        response.setMessageId(message.getMessageId());
        sendMessage(response);
    }

    private void handleHeartbeat(HeartbeatMessage message) {
        String requestWorkerId = message.getWorkerId(); // 原始请求中的 workerId

        // 处理心跳请求
        if (requestWorkerId != null) {
            // 这是来自 Worker 的心跳
            if (server.getServerType() == ClusterServer.ServerType.MASTER) {
                server.getRouteManager().updateHeartbeat(requestWorkerId);
                logger.debug("Updated heartbeat for worker: {}", requestWorkerId);
            } else if (server.getServerType() == ClusterServer.ServerType.WORKER) {
                // Worker 不应接收其他 Worker 的心跳，忽略
                logger.trace("Received worker heartbeat on worker node");
            }
        } else {
            // 客户端心跳
            logger.trace("Client heartbeat from {}", getClientAddress());
        }

        // 创建并发送响应
        HeartbeatResponse response = new HeartbeatResponse();
        response.setMessageId(message.getMessageId());

        // 关键修复：返回请求中的 workerId（如果存在）
        // 这样 Worker 就能验证响应
        response.setWorkerId(requestWorkerId);

        sendMessage(response);
    }
}